{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54915dac-fca2-410b-999b-710d50ded0f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from onediff.infer_compiler import oneflow_compile\n",
    "from onediff.schedulers import EulerDiscreteScheduler\n",
    "from diffusers import StableDiffusionPipeline\n",
    "import oneflow as flow\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8f6655b-a289-4da9-ba24-ea491fa2610d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"runwayml/stable-diffusion-v1-5\"\n",
    "prompt = \"photo of a cat\"\n",
    "height = 512\n",
    "width = 512\n",
    "steps = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0b75cd0-2e0f-4d63-b470-95f50c1d6e60",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder=\"scheduler\")\n",
    "pipe = StableDiffusionPipeline.from_pretrained(\n",
    "    model_id,\n",
    "    scheduler=scheduler,\n",
    "    torch_dtype=torch.float16,\n",
    "    safety_checker=None,\n",
    ")\n",
    "pipe = pipe.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "686526cc",
   "metadata": {},
   "source": [
    "### Now if you don't want to compile the model each time, you can cache it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad264c64",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.unet = oneflow_compile(\n",
    "    pipe.unet,\n",
    "    options={\n",
    "        \"graph_file\": \"/stable-diffusion-v1-5_graph\",\n",
    "        \"graph_file_device\": \"cuda\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9e81e15",
   "metadata": {},
   "source": [
    "So now when you run your first image generation, it will take a while to compile the model, but then it will be cached at `/stable-diffusion-v1-5_graph`.\n",
    "Next time you run image generation, it will be much faster."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe1ec1cf",
   "metadata": {},
   "source": [
    "### Or just run without caching:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f1b42a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.unet = oneflow_compile(pipe.unet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaec58c1-7262-48c6-866e-51fe8bc9c295",
   "metadata": {},
   "outputs": [],
   "source": [
    "images = pipe(\n",
    "    prompt, height=height, width=width, num_inference_steps=steps\n",
    ").images\n",
    "image = images[0]\n",
    "image"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
